#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
录制线程模块
"""

import os
import time
import mss
import numpy as np
import cv2
from PyQt5.QtCore import QThread, pyqtSignal, QTimer
from PyQt5.QtMultimedia import QAudioRecorder, QAudioEncoderSettings

class RecordingTimer(QThread):
    """录制时间计时器线程"""
    update_time = pyqtSignal(int)  # 发送已录制的秒数
    
    def __init__(self):
        super().__init__()
        self.running = False
        self.seconds = 0
    
    def run(self):
        self.running = True
        self.seconds = 0
        while self.running:
            self.seconds += 1
            self.update_time.emit(self.seconds)
            time.sleep(1)
    
    def stop(self):
        self.running = False
        self.wait()

class ScreenRecorderThread(QThread):
    """屏幕录制线程类"""
    # 定义信号
    update_frame = pyqtSignal(np.ndarray)  # 用于更新预览帧
    recording_stopped = pyqtSignal()  # 用于通知录制已停止
    recording_error = pyqtSignal(str)  # 用于通知录制错误
    recording_progress = pyqtSignal(int)  # 用于更新录制进度

    def __init__(self, output_file, fps=15.0, preview=False, monitor_num=1, record_audio=False, region=None):
        super().__init__()
        self.output_file = output_file
        self.fps = fps
        self.preview = preview
        self.monitor_num = monitor_num - 1  # mss使用0为第一个显示器
        self.record_audio = record_audio
        self.region = region  # 录制区域，格式为 (left, top, width, height)
        self.running = False
        self.paused = False
        self.audio_recorder = None
        self.audio_file = None
        
        # 进度更新计时器
        self.progress_timer = None
        self.start_time = None
        self.total_frames = 0

    def setup_audio_recorder(self):
        """设置音频录制"""
        if not self.record_audio:
            return
            
        try:
            self.audio_file = os.path.splitext(self.output_file)[0] + ".wav"
            self.audio_recorder = QAudioRecorder()
            settings = QAudioEncoderSettings()
            settings.setCodec("audio/pcm")
            settings.setQuality(QAudioEncoderSettings.HighQuality)
            self.audio_recorder.setEncodingSettings(settings)
            self.audio_recorder.setOutputLocation(self.audio_file)
        except Exception as e:
            self.recording_error.emit(f"音频录制设置失败: {str(e)}")
            self.record_audio = False

    def run(self):
        """线程运行函数"""
        try:
            self.running = True
            self.paused = False
            self.start_time = time.time()
            self.total_frames = 0
            
            # 设置进度更新计时器
            self.progress_timer = QTimer()
            self.progress_timer.timeout.connect(self.update_progress)
            self.progress_timer.start(1000)  # 每秒更新一次
            
            # 设置音频录制
            if self.record_audio:
                self.setup_audio_recorder()
                if self.audio_recorder:
                    self.audio_recorder.record()
            
            # 获取屏幕大小和区域
            with mss.mss() as sct:
                try:
                    if self.region:
                        # 使用自定义区域
                        monitor = {
                            "left": self.region[0],
                            "top": self.region[1],
                            "width": self.region[2],
                            "height": self.region[3]
                        }
                    else:
                        # 使用整个显示器
                        try:
                            monitor = sct.monitors[self.monitor_num]
                        except IndexError:
                            self.recording_error.emit(f"显示器 {self.monitor_num + 1} 不存在，使用主显示器")
                            monitor = sct.monitors[0]  # 使用主显示器
                    
                    width = monitor['width']
                    height = monitor['height']
                    
                except Exception as e:
                    self.recording_error.emit(f"获取屏幕区域失败: {str(e)}")
                    return

                # 设置视频编码器和输出
                fourcc = cv2.VideoWriter_fourcc(*"mp4v")
                out = cv2.VideoWriter(self.output_file, fourcc, self.fps, (width, height))

                # 开始录制
                start_time = time.time()
                frame_count = 0

                while self.running:
                    if not self.paused:
                        # 捕获屏幕
                        screenshot = sct.grab(monitor)
                        # 转换为numpy数组
                        frame = np.array(screenshot)
                        # 转换颜色格式 (BGRA to BGR)
                        frame = cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR)
                        # 写入视频
                        out.write(frame)
                        frame_count += 1
                        self.total_frames += 1

                        # 发送信号更新预览
                        if self.preview and frame_count % 3 == 0:  # 每3帧更新一次预览，减少CPU负担
                            self.update_frame.emit(frame)

                        # 控制帧率
                        elapsed = time.time() - start_time
                        expected_frames = int(elapsed * self.fps)
                        if frame_count > expected_frames:
                            self.msleep(int((frame_count - expected_frames) * 1000 / self.fps))
                    else:
                        # 暂停状态，休眠一小段时间
                        self.msleep(100)

                # 释放资源
                out.release()
                
                # 停止音频录制
                if self.record_audio and self.audio_recorder:
                    self.audio_recorder.stop()
                    
                    # 合并音视频（如果有音频）
                    if os.path.exists(self.audio_file) and os.path.getsize(self.audio_file) > 0:
                        self.merge_audio_video()
                
                if self.progress_timer:
                    self.progress_timer.stop()
                    
                self.recording_stopped.emit()
                
        except Exception as e:
            self.recording_error.emit(f"录制过程中出错: {str(e)}")
            if self.progress_timer:
                self.progress_timer.stop()
            self.recording_stopped.emit()

    def update_progress(self):
        """更新录制进度"""
        if self.start_time:
            elapsed = int(time.time() - self.start_time)
            self.recording_progress.emit(elapsed)

    def merge_audio_video(self):
        """合并音频和视频文件"""
        try:
            # 这里可以使用ffmpeg合并音视频，但需要安装ffmpeg
            # 为简化依赖，这里只保留音频文件，让用户自行合并
            pass
        except Exception as e:
            self.recording_error.emit(f"合并音视频失败: {str(e)}")

    def pause(self):
        """暂停录制"""
        self.paused = True
    
    def resume(self):
        """恢复录制"""
        self.paused = False

    def stop(self):
        """停止录制"""
        self.running = False
        if self.audio_recorder:
            self.audio_recorder.stop()
        self.wait()
        
    def is_paused(self):
        """检查是否暂停状态"""
        return self.paused